import torch

from utils.visualize import visualize
import os
import pandas as pd
from utils.reranker import Reranker
import logging

# 获取 vllm 日志记录器
vllm_logger = logging.getLogger("vllm")

# 设置日志等级，例如设置为 WARNING
vllm_logger.setLevel(logging.WARNING)

# os.environ["CUDA_VISIBLE_DEVICES"] = "6,7"

# model_path = "model/Qwen/Qwen2-VL-2B-Instruct-AWQ"
model_path = "/home/hutu/workspace/qwen_server/model/qwen2-vl-7B-Instruct-AWQ"

print(torch.cuda.device_count())

/home/hutu/miniconda3/lib/python3.12/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html
  from .autonotebook import tqdm as notebook_tqdm
2025-01-10 16:31:00,805	INFO util.py:154 -- Missing packages: ['ipywidgets']. Run `pip install -U ipywidgets`, then restart the notebook server for rich notebook output.

2

reranker = Reranker(
    model_path, quantization="awq", tensor_parallel_size=2, max_model_len=10000
)

WARNING 01-10 16:31:12 config.py:428] awq quantization is not fully optimized yet. The speed can be slower than non-quantized models.
WARNING 01-10 16:31:12 config.py:503] To see benefits of async output processing, enable CUDA graph. Since, enforce-eager is enabled, async output processor cannot be used
WARNING 01-10 16:31:12 multiproc_gpu_executor.py:56] Reducing Torch parallelism from 44 threads to 1 to avoid unnecessary CPU contention. Set OMP_NUM_THREADS in the external environment to tune this value as needed.

huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...
To disable this warning, you can either:
	- Avoid using `tokenizers` before the fork if possible
	- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)

WARNING 01-10 16:31:18 custom_all_reduce.py:143] Custom allreduce is disabled because your platform lacks GPU P2P capability or P2P test failed. To silence this warning, specify disable_custom_all_reduce=True explicitly.
(VllmWorkerProcess pid=1359974) WARNING 01-10 16:31:18 custom_all_reduce.py:143] Custom allreduce is disabled because your platform lacks GPU P2P capability or P2P test failed. To silence this warning, specify disable_custom_all_reduce=True explicitly.

Loading safetensors checkpoint shards:   0% Completed | 0/2 [00:00<?, ?it/s]
Loading safetensors checkpoint shards:  50% Completed | 1/2 [00:01<00:01,  1.04s/it]
Loading safetensors checkpoint shards: 100% Completed | 2/2 [00:01<00:00,  1.29it/s]
Loading safetensors checkpoint shards: 100% Completed | 2/2 [00:01<00:00,  1.22it/s]

# # read parquet file
# df = pd.read_parquet("data/all_hw_images_cleaned.parquet")
# df.head(1)

df = pd.read_json("data/shopee/marqo.json", lines=True)

# generate labels
label2posting_id = df.groupby("label_group").posting_id.agg("unique").to_dict()
df["label"] = df.label_group.map(label2posting_id)

df["image_path"] = df["image"].apply(
    lambda x: os.path.join("data/shopee/train_images", x)
)
df.head(1)

# iterate over the rows of the dataframe
for i, row in df.iterrows():
    query_text = row["title"]
    query_image_path = row["image_path"]
    preds_ids = row["pred"]
    preds_df = df[df["posting_id"].isin(preds_ids)]

    preds_ids = preds_df["posting_id"].tolist()
    preds_text = preds_df["title"].tolist()
    preds_image_path = preds_df["image_path"].tolist()

    reranked_ids = reranker.rerank(
        query_text,
        preds_text,
        query_image_path,
        preds_image_path,
        preds_ids,
        use_tqdm=True,
    )
    gt_ids = row["label"]

    reranked_df = df[df["posting_id"].isin(reranked_ids)]
    reranked_ids = reranked_df["posting_id"].tolist()
    reranked_text = reranked_df["title"].tolist()
    reranked_image_path = reranked_df["image_path"].tolist()
    visualize(query_text, reranked_text, query_image_path, reranked_image_path)

    gt_df = df[df["posting_id"].isin(gt_ids)]
    gt_ids = gt_df["posting_id"].tolist()
    gt_text = gt_df["title"].tolist()
    gt_image_path = gt_df["image_path"].tolist()
    visualize(query_text, gt_text, query_image_path, gt_image_path)
    if i == 20:
        break

Processed prompts: 100%|██████████| 15/15 [00:22<00:00,  1.49s/it, est. speed input: 1621.74 toks/s, output: 0.67 toks/s]

Processed prompts: 100%|██████████| 15/15 [00:20<00:00,  1.35s/it, est. speed input: 1702.82 toks/s, output: 0.74 toks/s]

Processed prompts: 100%|██████████| 15/15 [00:12<00:00,  1.18it/s, est. speed input: 1805.99 toks/s, output: 1.18 toks/s]

Processed prompts: 100%|██████████| 15/15 [00:14<00:00,  1.04it/s, est. speed input: 1782.88 toks/s, output: 1.04 toks/s]

Processed prompts: 100%|██████████| 15/15 [00:13<00:00,  1.12it/s, est. speed input: 1834.57 toks/s, output: 1.12 toks/s]

Processed prompts: 100%|██████████| 15/15 [00:12<00:00,  1.23it/s, est. speed input: 1828.42 toks/s, output: 1.23 toks/s]

Processed prompts: 100%|██████████| 15/15 [00:11<00:00,  1.32it/s, est. speed input: 1840.19 toks/s, output: 1.32 toks/s]

Processed prompts: 100%|██████████| 15/15 [00:10<00:00,  1.43it/s, est. speed input: 1897.37 toks/s, output: 1.43 toks/s]

Processed prompts: 100%|██████████| 15/15 [00:08<00:00,  1.85it/s, est. speed input: 1919.84 toks/s, output: 1.85 toks/s]

Processed prompts: 100%|██████████| 15/15 [00:12<00:00,  1.21it/s, est. speed input: 1874.21 toks/s, output: 1.21 toks/s]

Oversea Dataset¶

Shopee Dataset¶